#!/usr/bin/env python
# coding=utf8

import random
import math
import show
import simulation
import numpy as np
from LSTMModule import CLASS_NUM


def findTranPoint(datalist, maxpoint=10, span=300):
    for one in datalist:
        one["tran_mark"] = int(CLASS_NUM / 2)
    for i in range(span, len(datalist) - span):
        testdata = datalist[i - span//500:i + span]
        testdata = sorted(testdata, key=lambda a: a["DI"])
        tp = datalist[i]
        if tp in testdata[: maxpoint]:
            tp["tran_mark"] = 0
        elif tp in testdata[maxpoint: maxpoint * 3]:
            tp["tran_mark"] = 1
        elif tp in testdata[maxpoint * 3: maxpoint * 6]:
            tp["tran_mark"] = 2
        elif tp in testdata[-6 * maxpoint: -3 * maxpoint]:
            tp["tran_mark"] = 4
        elif tp in testdata[-3 * maxpoint: -maxpoint]:
            tp["tran_mark"] = 5
        elif tp in testdata[-maxpoint:]:
            tp["tran_mark"] = 6


def fetchTranlist(datalist, fun, nosign=0):
    srclist = []
    reslist = []
    for i in range(100, len(datalist)):
        one = datalist[i]
        tl = fun(datalist[:i + 1])
        drop = False
        for oned in tl:
            if math.isnan(oned):
                drop = True
                break
        if drop:
            continue
        if "tran_mark" in one:
            passdata = False
            for t in tl:
                if math.isnan(t):
                    passdata = True
                    break
            if passdata: continue
            srclist.append(tl)
            reslist.append(one["tran_mark"])
        elif random.random() < nosign:
            srclist.append(tl)
            reslist.append(0)
    return srclist, reslist


def splitTrainTest(inputs, outputs, trainpick=0.7):
    train_output = []
    test_output = []
    train_input = []
    test_input = []
    for i in range(len(inputs)):
        hasnan = False
        for ck in inputs[i]:
            if math.isnan(ck):
                hasnan = True
                break
        if hasnan:
            continue
        if random.random() < trainpick:
            train_input.append(inputs[i])
            train_output.append(outputs[i])
        else:
            test_input.append(inputs[i])
            test_output.append(outputs[i])
    return train_input, train_output, test_input, test_output


def exDrawLine(axs, spandata):
    show.DrawTranPoint(axs[0], spandata)


if __name__ == "__main__":
    datalist = fetchdata.loaddate("data/golddata.gz")
    findTranPoint(datalist)

    def getone(one):
        return one["macd"][0], one["macd"][1], one["macd"][2], \
               one["sar"] - one["DI"], one["rsi"], one["bband"][0] - one["High"], one["bband"][2] - one["Low"]

    srclist, reslist = fetchTranlist(datalist, getone)
    train_input, train_output, test_input, test_output = splitTrainTest(srclist, reslist)
    fetchdata.savedata(datalist, "data/golddata.gz")
    show.Show(datalist[-500:], exDrawLine)
